""""
通过 pip install scikit-learn安装依赖
"""
from sklearn.linear_model import LinearRegression
import numpy as np

from vnpy_portfoliostrategy import (
    StrategyTemplate,
    ArrayManager,
    BarData,
    Direction
)


def calculate_score(data: np.ndarray) -> float:
    """计算强弱得分"""
    # 执行回归
    x: np.ndarray = np.arange(1, len(data) + 1).reshape(-1, 1)
    y: np.ndarray = data / data[0]
    reg: LinearRegression = LinearRegression().fit(x, y)

    # 返回得分
    slope: float = reg.coef_[0]
    r2: float = reg.score(x, y)
    return slope * r2


class EtfRotationStrategy(StrategyTemplate):
    """ETF轮动策略"""

    author: str = "CZL"

    regression_window: int = 25         # 线性回归窗口
    holding_size = 2                    # 最大持仓数
    fixed_capital: int = 1_000_000      # 固定持仓市值

    parameters = [
        "regression_window",
        "holding_size"
    ]

    def on_init(self) -> None:
        """策略初始化"""
        # 确保缓存数据足够回归计算
        size: int = self.regression_window + 1

        # 创建每个合约的时序数据容器
        self.ams: dict[str, ArrayManager] = {}

        for vt_symbol in self.vt_symbols:
            self.ams[vt_symbol] = ArrayManager(size)

        # 创建持仓合约名称的列表
        self.holding_symbols = []

        self.write_log("策略初始化")

    def on_start(self) -> None:
        """策略启动"""
        self.write_log("策略启动")

    def on_stop(self) -> None:
        """策略停止"""
        self.write_log("策略停止")

    def on_bars(self, bars: dict[str, BarData]) -> None:
        """K线切片推送"""
        # 更新K线到时序容器
        for vt_symbol, bar in bars.items():
            am: ArrayManager = self.ams[vt_symbol]
            am.update_bar(bar)

        # 计算每只ETF的分数
        score_data: dict[str, float] = {}

        for vt_symbol, bar in bars.items():
            am: ArrayManager = self.ams[vt_symbol]
            if not am.inited:
                return

            data: np.array = am.close[-self.regression_window:]
            score_data[vt_symbol] = calculate_score(data)

        # 重置所有合约目标
        for vt_symbol in self.vt_symbols:
            self.set_target(vt_symbol, 0)

        # 选出得分领先的ETF
        top_ranked = sorted(score_data, 
                            key=lambda x: score_data[x],
                            reverse=True)[:self.holding_size]

        # 如果得分领先的ETF不再属于现有持仓
        if set(top_ranked) != set(self.holding_symbols):
            # 更新目标持仓列表
            self.holding_symbols = top_ranked
            for security in self.holding_symbols:
                price: float = bars[security].close_price
                volume: int = 100 * int((self.fixed_capital / self.holding_size) / (price * 100))
                self.set_target(security, volume)

            # 根据设置好的目标仓位进行交易
            self.rebalance_portfolio(bars)

        # 推送UI更新
        self.put_event()

    def calculate_price(self, vt_symbol, direction, reference):
        if direction == Direction.LONG:
            price: float = round(reference * 1.05, 3)
        else:
            price = round(reference * 0.95, 3)

        return price
